#!/usr/bin/env python3

DBName = 'ping'
Host = 'baidu.com'

import time
import logging
import socket
from functools import partial

from tornado import ioloop
from pymongo import MongoClient
DB = None

import icmplib
from nicelogger import enable_pretty_logging
enable_pretty_logging()

def setup_sync():
  global DB, Host

  DB = MongoClient()[DBName]
  logging.info('MongoDB connected')
  if 'ping' not in DB.collection_names():
    DB.create_collection('ping', capped=True, size=1024 * 1024 * 256)
  DB.ping.ensure_index('t')
  logging.info('database setup done')

  ip = socket.gethostbyname(Host)
  logging.info('Host %s resolved to %s', Host, ip)
  Host = ip

class Pinger:
  seq = 0
  wait_timeout = 30 # this long or never
  payload = b'L' * 56
  last_received = 0
  sock = None

  def __init__(self, addr, io_loop=None):
    self.addr = addr
    self.io_loop = io_loop or ioloop.IOLoop().current()
    self.new_sock()

    self.flying = {}
    self.maylost = set()
    self.p = ioloop.PeriodicCallback(self.ping_host, 1000, io_loop=io_loop)
    self.p.start()
    self.ping_host()

  def new_sock(self):
    if self.sock:
      self.io_loop.remove_handler(self.sock.fileno())
    self.sock = icmplib.socket()
    self.io_loop.add_handler(self.sock.fileno(), self.pong_received, ioloop.IOLoop.READ)
    try:
      self.sock.connect((self.addr, 0))
    except OSError as e:
      if e.errno == 101:
        self.io_loop.remove_handler(self.sock.fileno())
        logging.error('Network unreachable. Will try again later.')
        self.sock = None
      else:
        raise

  def ping_host(self):
    self.seq = (self.seq + 1) & 0x7fff
    data = icmplib.pack_packet(self.seq, self.payload)
    t = time.time()
    self.flying[self.seq] = (t, self.io_loop.add_timeout(
      t + self.wait_timeout,
      partial(self.pong_never, self.seq),
    ))
    try:
      self.sock.send(data)
    except AttributeError:
      self.new_sock()
    except OSError as e:
      logging.error('Error while sending data. Creating new socket.')
      if e.errno == 22: # Invalid argument
        self.new_sock()
      else:
        raise

  def pong_received(self, sockfd, event):
    data = self.sock.recv(1024)
    seq = icmplib.parse_packet(data)[0]
    if (seq - self.last_received) % 0x7ffff > 1:
        for s in range(self.last_received+1, seq):
            if s > 0x7fff:
                s -= 0x7fff
            try:
                logging.warn('PONG %4d    NOT SEEN from %.3f', s, self.flying[s][0])
                self.maylost.add(s)
            except KeyError:
                # already printed lost
                pass
    self.last_received = seq
    try:
        t, timeout = self.flying[seq]
    except KeyError:
        # finally arrived, but we thought it was lost
        logging.info('PONG %4d finally arrived', seq)
        return
    self.io_loop.remove_timeout(timeout)
    interval = (time.time() - t) * 1000
    self.save_ping_result(t, interval)
    if seq in self.maylost:
        fmt = 'PONG %4d %9.3fms from %.3f (Out-of-Order)'
        self.maylost.remove(seq)
    else:
        fmt = 'PONG %4d %9.3fms from %.3f'
    logging.debug(fmt, seq, interval, t)

  def save_ping_result(self, t, interval):
    DB.ping.save({
      't': t,
      'i': interval,
    })

  def pong_never(self, seq):
    t, _ = self.flying[seq]
    if seq in self.maylost:
        self.maylost.remove(seq)
    else:
        logging.warn('PONG %4d        LOST from %.3f', seq, t)
    self.save_ping_result(t, float('inf'))
    del self.flying[seq]

  def __del__(self):
    self.p.stop()

def main():
  setup_sync()
  loop = ioloop.IOLoop().instance()
  Pinger(Host)
  loop.start()

if __name__ == '__main__':
  try:
    import setproctitle
    setproctitle.setproctitle('pingd')
    del setproctitle
  except ImportError:
    pass

  try:
    main()
  except KeyboardInterrupt:
    print()
